import os
import json
import requests
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()  # read local .env file

class Agent:
    def __init__(self, prompt):
        """
        初始化 LanguageModelAgent 实例。

        :param api_key: OpenAI API 密钥
        :param base_url: OpenAI API 基础 URL
        :param prompt: 系统提示，指导模型的行为
        """
        self.model = OpenAI(api_key=os.getenv("API_KEY"),
                            base_url=os.getenv("BASE_URL"))

        self.prompt = prompt

        self.messages = [{"role": "system", "content": prompt}]

    def chat(self, user_input, context="", stream=True, temperature=0.6):
        """
        处理用户输入并生成响应。

        :param user_input: 用户输入
        :param context: 上下文信息，默认为空字符串
        :return: 模型生成的响应
        """
        user_message = f"{user_input}."

        if context:
            user_message = f"已知\n:{context}" + "\n" + user_message

        self.messages.append({"role":"user", "content":user_message})

        response = self.model.chat.completions.create(
            #model="deepseek-chat",
            model = "deepseek-reasoner",
            messages=self.messages,
            temperature=temperature,
            stream=stream)

        if stream:
            return response
        else:
            self.messages.append(response.choices[0].message)
            return response.choices[0].message.content

    def collect_message(self, gen_txt):
        '''当使用 st.stream_write 时把输出结果保存在 messages 列表中'''
        self.messages.append({"role": "assistant", "content": gen_txt})



class AgentLocal:
    def __init__(self, prompt):
        """
        初始化 LanguageModelAgent 实例。

        :param prompt: 系统提示，指导模型的行为
        """

        self.prompt = prompt

        self.messages = [{"role": "system", "content": prompt}]
 

    def chat(self, user_input, context="", stream=True, temperature=0.6):
        """
        处理用户输入并生成响应。

        :param user_input: 用户输入
        :param context: 上下文信息，默认为空字符串
        :return: 模型生成的响应
        """
        user_message = f"{user_input}."

        if context:
            user_message = f"已知\n:{context}" + "\n" + user_message

        self.messages.append({"role":"user", "content":user_message})

        # 修改为使用本地部署的 Ollama deepseek
        url = "http://localhost:11434/v1/chat/completions"
        headers = {
            "Content-Type": "application/json"
        }
        data = {
            "model": "deepseek-r1:14b",
            "messages": self.messages
        }

        response = requests.post(url, headers=headers, json=data)

        if response.status_code == 200:
            response_data = response.json()
            self.messages.append(response_data['choices'][0]['message'])
            return response_data['choices'][0]['message']['content']
        else:
            print(f"请求失败，状态码：{response.status_code}")
            return None

    def collect_message(self, gen_txt):
        '''当使用 st.stream_write 时把输出结果保存在 messages 列表中'''
        self.messages.append({"role": "assistant", "content": gen_txt})




def rm_mkdown(gen_txt):
    code = gen_txt.replace("```json", "")
    code = code.replace("```python", "")
    code = code.replace("```", "")
    code = code.replace("\n", "")
    return code


if __name__ == "__main__":
    agent = Agent("你是一个非常厉害的科研助理。你会回答我的物理问题。")
    print(agent.chat("你好"))
